import cv2
import numpy as np

# 定义鼠标点击事件
points_left = [719, 286]  # 存储左图像点击点
points_right = [567, 285]  # 存储右图像点击点
points_left = [1413, 1027]  # 存储左图像点击点
points_right = [1326, 1027]  # 存储右图像点击点
# points_left = []  # 存储左图像点击点
# points_right = []  # 存储右图像点击点


# 鼠标回调函数
def click_event(event, x, y, flags, param):
    global points_left, points_right

    if event == cv2.EVENT_LBUTTONDOWN:
        if len(points_left) < len(points_right):
            points_left.append((x, y))
            print(f"左图像点击点：{(x, y)}")
        else:
            points_right.append((x, y))
            print(f"右图像点击点：{(x, y)}")

        # 标记点击的位置
        cv2.circle(param, (x, y), 5, (0, 0, 255), -1)
        cv2.imshow("Image", param)


# # 左相机的内参矩阵 (fx, fy, u0, v0)
# K1 = np.array([[9060.5738644625944, 0, 770.0566478118638],
#                [0, 9083.2694501407805, 609.20841979353941],
#                [0, 0, 1]])
#
# # 右相机的内参矩阵 (fx, fy, u0, v0)
# K2 = np.array([[9084.6147664638211, 0, 740.41002940402359],
#                [0, 9107.1443164906887, 618.90885533722155],
#                [0, 0, 1]])
#
# # 左相机畸变系数 (k1, k2, p1, p2, k3)
# D1 = np.array([0.09784518980218336, 5.2606551678446252, 0.0, 0.0, 0.0])
#
# # 右相机畸变系数 (k1, k2, p1, p2, k3)
# D2 = np.array([-0.056758282671412398, 11.950457152339208, 0.0, 0.0, 0.0])
#
# # 旋转矩阵 R
# R = np.array([[0.99999449726182721, -0.000068181879709459687, 0.0033167449851725151],
#               [0.00007190546024448005, 0.9999993673499662, -0.0011225547969415993],
#               [-0.0033166663489375557, 0.0011227871118911395, 0.99999386951792424]])
#
# # 平移向量 T
# T = np.array([-200.70787924046954, 0.40572409096883688, 15.226810631241822])
#
#
# # 左相机的内参矩阵 (fx, fy, u0, v0)
# K1 = np.array([[9.1025453266397271e+03, 0, 9.8143301848885960e+02],
#                [0, 9.0954119671537210e+03, 6.0517957568967927e+02],
#                [0, 0, 1]])
#
# # 右相机的内参矩阵 (fx, fy, u0, v0)
# K2 = np.array([[9.1044184325803071e+03, 0, 9.7187568822509047e+02],
#                [0, 9.0977242381526758e+03, 6.0841419215181554e+02],
#                [0, 0, 1]])
#
# # 左相机畸变系数 (k1, k2, p1, p2, k3)
# D1 = np.array([1.5584032577386640e-02, 7.6614106768244694e-01, 0.0, 0.0, 0.0])
#
# # 右相机畸变系数 (k1, k2, p1, p2, k3)
# D2 = np.array([3.5368566264021529e-03, 1.0593839296711431e+00, 0.0, 0.0, 0.0])
#
# # 旋转矩阵 R
# R = np.array([[9.9999946263303441e-01, -9.2023173997801200e-06, 1.0366527672624633e-03],
#               [9.5660266554949309e-06, 9.9999993840805479e-01, -3.5084523337578795e-04],
#               [-1.0366494748238074e-03, 3.5085496149115297e-04, 9.9999940112915187e-01]])
#
# # 平移向量 T
# T = np.array([-1.9978096828715928e+02, -2.9008954988253656e-02, 1.4517684084484477e+00])


# 左相机的内参矩阵 (fx, fy, u0, v0)
K1 = np.array([[9.1565453266397271e+03, 0, 960],
               [0, 9.1494119671537210e+03, 600],
               [0, 0, 1]])

# 右相机的内参矩阵 (fx, fy, u0, v0)
K2 = np.array([[9.1644184325803071e+03, 0, 960],
               [0, 9.1517242381526758e+03, 600],
               [0, 0, 1]])

# 左相机畸变系数 (k1, k2, p1, p2, k3)
D1 = np.array([0.0, 0.0, 0.0, 0.0, 0.0])

# 右相机畸变系数 (k1, k2, p1, p2, k3)
D2 = np.array([0.0, 0.0, 0.0, 0.0, 0.0])

# 旋转矩阵 R
R = np.array([[1.0, 0.0, 0.0],
              [0.0, 1.0, 0.0],
              [0.0, 0.0, 1.0]])

# 平移向量 T
T = np.array([-200.0, 0.0, 0.0])


# 计算立体标定的重投影矩阵
R1, R2, P1, P2, Q, roi1, roi2 = cv2.stereoRectify(K1, D1, K2, D2, (1920, 1200), R, T)

# 左右图像的读取
left_image = cv2.imread(r"D:\Desktop\RoboMaster-Unity\Screenshots\left.png")
right_image = cv2.imread(r"D:\Desktop\RoboMaster-Unity\Screenshots\right.png")

# 校正图像
map1_left, map2_left = cv2.initUndistortRectifyMap(K1, D1, R1, P1, (1920, 1200), cv2.CV_32FC1)
map1_right, map2_right = cv2.initUndistortRectifyMap(K2, D2, R2, P2, (1920, 1200), cv2.CV_32FC1)

left_rectified = cv2.remap(left_image, map1_left, map2_left, cv2.INTER_LINEAR)
right_rectified = cv2.remap(right_image, map1_right, map2_right, cv2.INTER_LINEAR)

# # 显示左图像并获取点击点
# cv2.namedWindow("Left Image", 0)
# cv2.imshow("Left Image", left_rectified)
# cv2.setMouseCallback("Left Image", click_event, left_rectified)
# cv2.waitKey(0)
#
# # 显示右图像并获取点击点
# cv2.namedWindow("Right Image", 0)
# cv2.imshow("Right Image", right_rectified)
# cv2.setMouseCallback("Right Image", click_event, right_rectified)
# cv2.waitKey(0)

# # 确保左右目都点击完毕
# if len(points_left) == len(points_right):
#     print("获取到对应的点对")
#
#     # 计算三维坐标
#     for i in range(len(points_left)):
#         # 获取左图和右图中的对应点
#         p_left = np.array([points_left[i][0], points_left[i][1], 1])  # 归一化坐标
#         p_right = np.array([points_right[i][0], points_right[i][1], 1])  # 归一化坐标
#
#         # 通过立体校正和三角化来计算三维坐标
#         points_3d = cv2.triangulatePoints(K1, K2, p_left[:2], p_right[:2])
#         points_3d /= points_3d[3]  # 齐次坐标转换为笛卡尔坐标
#         print(f"三维坐标：{points_3d[:3]}")


def get_3d_position_by_matrix(Q, p_l, p_r):
    """
    通过立体校正矩阵 Q 和左右目图像的点 p_l, p_r 计算三维坐标。

    :param Q: 立体校正矩阵 (4x4)
    :param p_l: 左图像中的点 (x, y)
    :param p_r: 右图像中的点 (x, y)
    :return: 三维坐标 (X, Y, Z)
    """
    # 创建 armorMtx 向量 (4x1)，包括 p_l.x, p_l.y, p_l.x - p_r.x, 1
    armor_mtx = np.array([[p_l[0]], [p_l[1]], [p_l[0] - p_r[0]], [1.0]])

    print(Q)
    print(armor_mtx)

    # 使用 Q 进行矩阵乘法
    cal_q = np.dot(Q, armor_mtx)

    # 获取三维坐标（X, Y, Z）和齐次坐标 W
    X, Y, Z, W = cal_q.flatten()

    # 返回三维坐标
    return X / W, Y / W, Z / W

X, Y, Z = get_3d_position_by_matrix(Q, points_left, points_right)

print(f"三维坐标: X={X}, Y={Y}, Z={Z}")


def draw_h_concat_line(img_l, img_r):
    """
    拼接左右目图像并在竖直方向绘制直线。

    :param img_l: 左图像
    :param img_r: 右图像
    :return: 拼接后的图像，包含绘制的直线
    """
    kInterval = 25  # 所绘直线竖直方向的间隔
    kColors = [
        (75, 72, 242),  # 颜色1
        (254, 164, 67),  # 颜色2
        (144, 203, 251),  # 颜色3
        (255, 110, 29),  # 颜色4
        (24, 95, 141),  # 颜色5
        (178, 114, 20)  # 颜色6
    ]

    # 拼接左右图像
    output = np.hstack((img_l, img_r))

    height, width = output.shape[:2]

    # 绘制直线
    for i in range(kInterval, height, kInterval):
        color = kColors[(i // kInterval) % 6]  # 循环使用颜色
        cv2.line(output, (0, i), (width, i), color, 1)  # 绘制水平线

    return output



# 绘制拼接并显示
output = draw_h_concat_line(left_rectified, right_rectified)
cv2.namedWindow("Output Image", 0)
cv2.imshow("Output Image", output)
cv2.waitKey(0)

cv2.destroyAllWindows()
